package com.mdframe.ai.agent.config;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.CommandLineRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.Resource;

/**
 * @date 2024/10/12
 */
@Slf4j
@Configuration
public class AssistantConfig {
    
    
    @Value("classpath:rag/controller.txt") Resource controllerResource;
    @Value("classpath:rag/entity.txt") Resource entityResource;
    @Value("classpath:rag/mapper.txt") Resource mapperResource;
    @Value("classpath:rag/service.txt") Resource serviceResource;
    @Value("classpath:rag/serviceImpl.txt") Resource serviceImplResource;
    
    @Bean
    CommandLineRunner ingestTermOfServiceToVectorStore(EmbeddingModel embeddingModel, VectorStore vectorStore) {

        return args -> {
            // Ingest the document into the vector store
            vectorStore.write(new TokenTextSplitter().transform(new TextReader(controllerResource).read()));
            log.info("向量初始化完成");
        };
    }

    @Bean
    public VectorStore vectorStore(EmbeddingModel embeddingModel) {
        return new SimpleVectorStore(embeddingModel);
    }
    
    @Bean
    public ChatMemory chatMemory() {
        return new InMemoryChatMemory();
    }
}
